package sgxnb

/*
#cgo CFLAGS: -I/opt/intel/sgxsdk/include
#cgo LDFLAGS: -L/opt/intel/sgxsdk/lib64 -lsgx_urts -lsgx_uae_service
#include "sgx_urts.h"
*/
import "C"

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"unsafe"
)

type Enclave struct {
	ID     C.sgx_enclave_id_t
	Handle unsafe.Pointer
}

// 创建enclave
func CreateEnclave() (*Enclave, error) {
	var enclave Enclave
	var ret C.int

	// 创建enclave
	ret = C.sgxn_create_enclave(C.CString("enclave.signed.so"), C.int(len("enclave.signed.so")), &enclave.ID, &enclave.Handle)
	if ret != C.SGX_SUCCESS {
		return nil, fmt.Errorf("Failed to create enclave: %d", ret)
	}

	return &enclave, nil
}

// 销毁enclave
func DestroyEnclave(enclave *Enclave) error {
	var ret C.int

	// 销毁enclave
	ret = C.sgxn_destroy_enclave(enclave.ID)
	if ret != C.SGX_SUCCESS {
		return fmt.Errorf("Failed to destroy enclave: %d", ret)
	}

	return nil
}

// 在SGX里面加密数据
func EncryptSGX(plaintext []byte) ([]byte, error) {
	// 创建enclave
	enclave, err := CreateEnclave()
	if err != nil {
		return nil, err
	}
	defer DestroyEnclave(enclave)

	// 加密数据
	var sgxCiphertextAES []byte
	ret, err := C.sgxn_encrypt(enclave.ID, (*C.uchar)(&plaintext[0]), C.uint(len(plaintext)), (*C.uchar)(&sgxCiphertextAES[0]), C.uint(len(sgxCiphertextAES)))
	if ret != C.SGX_SUCCESS {
		return nil, fmt.Errorf("Failed to encrypt data: %d", ret)
	}

	// 加密AES密钥
	clientPublicKey, err := LoadPublicKey("clientB.pub")
	if err != nil {
		return nil, err
	}
	encryptedSgxCiphertextAES, err := EncryptRSAnb(sgxCiphertextAES, clientPublicKey)
	if err != nil {
		return nil, err
	}

	return encryptedSgxCiphertextAES, nil
}

// 加载公钥
func LoadPublicKey(path string) (*rsa.PublicKey, error) {
	// 读取公钥文件
	publicKeyFile, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}

	// 解码PEM格式的公钥
	block, _ := pem.Decode(publicKeyFile)
	if block == nil {
		return nil, errors.New("Failed to decode PEM block containing public key")
	}

	// 解析DER编码的公钥
	publicKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}

	// 转换为RSA公钥
	rsaPublicKey, ok := publicKey.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("Failed to convert public key to RSA")
	}

	return rsaPublicKey, nil
}

// 加密数据
func EncryptRSAnb(plaintext []byte, publicKey *rsa.PublicKey) ([]byte, error) {
	// 使用PKCS1v15填充方案加密数据
	ciphertext, err := rsa.EncryptPKCS1v15(rand.Reader, publicKey, plaintext)
	if err != nil {
		return nil, err
	}

	return ciphertext, nil
}
